pub mod tree;
pub mod error;

use std::{io::Cursor};

pub use tree::*;
pub use error::*;

use crate::{byte::*, fse::{FseDistribution, FseTable}};

#[derive(Debug, Clone, Copy)]
pub struct HuffmanEntry {
    pub(crate) literal:u8,
    pub(crate) weight:u8,
    pub(crate) number_of_bits:u8,
}

#[derive(Clone)]
pub struct HuffmanTable {
    max_bits:u8,
    huffman_tree:Box<HuffmanTree>,
}

impl HuffmanTable {
    pub fn serialize_into<W:ByteWrite>(&self, _writer:&mut ByteWriter<W>) -> HuffmanResult<()> {
        todo!()
    }

    pub fn max_bits(&self) -> u8 {
        self.max_bits
    }

    pub fn deserialize_from(raw_data:&[u8]) -> HuffmanResult<(Self, usize)> {
        let header= raw_data[0];
        let mut read_bytes_num = 1;
        let mut weight = Vec::new();

        if header < 128 {
            let compressed_size = header;
            read_bytes_num += compressed_size as usize;
            let raw_data = &raw_data[1..1 + compressed_size as usize];
            let mut total_len = raw_data.len() as u32;
            let mut reader = ByteReader::new(Cursor::new(raw_data));
            let distribution = FseDistribution::deserialize_from(&mut total_len, &mut reader)?;
            let fse_table:FseTable = distribution.into();

            let mut weight_stream = raw_data[raw_data.len() - total_len as usize..].to_vec();
            weight_stream.reverse();
            let weight_stream_bits_num = weight_stream.len() * 8;

            let mut byte_reader = ByteReader::new(Cursor::new(weight_stream));
            let mut bit_reader = BitReader::new(&mut byte_reader);
            bit_reader.read_reverse_padding()?;

            let mut even_cell = fse_table.first_cell(&mut bit_reader)?;
            let mut odd_cell = fse_table.first_cell(&mut bit_reader)?;

            loop {
                let even_symbol = even_cell.symbol;
                let odd_symbol = odd_cell.symbol;

                weight.push(even_symbol);
                weight.push(odd_symbol);

                let remain_bits = weight_stream_bits_num - bit_reader.bits_num();
                if remain_bits < even_cell.num_bits as usize {
                    let remain_num:u8 = bit_reader.read_bits(remain_bits)?;
                    assert_eq!(remain_num, 0);
                    break;
                } else if remain_bits < (even_cell.num_bits + odd_cell.num_bits) as usize {
                    let even_final_cell = fse_table.next_cell(&mut bit_reader, even_cell)?;
                    weight.push(even_final_cell.symbol);
                    let remain_num:u8 = bit_reader.read_bits(remain_bits - even_cell.num_bits as usize)?;
                    assert_eq!(remain_num, 0);
                    break;
                }

                even_cell = fse_table.next_cell(&mut bit_reader, even_cell)?;
                odd_cell = fse_table.next_cell(&mut bit_reader, odd_cell)?;
            }
            assert!(!bit_reader.bits_remaining());
        } else {
            let number_of_symbols = header - 127;
            let need_bytes_num = (number_of_symbols as usize + 1)/2;
            for byte in &raw_data[1..1 + need_bytes_num] {
                weight.push(byte >> 4);
                weight.push(byte & 0xf);
            }
            if number_of_symbols % 2 == 1 { weight.pop(); }
            read_bytes_num += need_bytes_num;
        };
        let mut prefix_weight_sum:u16 = 0;
        weight
            .iter()
            .filter(|w|**w!=0)
            .for_each(|w|{prefix_weight_sum += 1 << (*w - 1);});
        let max_bits = (u16::BITS - prefix_weight_sum.leading_zeros()) as u8;
        let total_weight_sum = 1 << max_bits;
        let last_weight_power = total_weight_sum - prefix_weight_sum;
        assert!(last_weight_power.is_power_of_two());
        let last_weight = u16::BITS - last_weight_power.leading_zeros();
        weight.push(last_weight as u8);

        let huffman_tree = Self::build_from_wight(max_bits, weight).ok_or(HuffmanError::Tree)?;
        Ok((Self { max_bits, huffman_tree }, read_bytes_num))
    }

    pub fn build_from_wight(max_bits:u8, weight:Vec<u8>) -> Option<Box<HuffmanTree>> {
        let mut literal_order:Vec<_> = weight
            .iter()
            .enumerate()
            .map(|(i, w)|
                HuffmanEntry{ literal:i as u8, weight:*w, number_of_bits:Self::weight2bits(max_bits, *w) })
            .collect();
        literal_order.sort_by(|a,b|
            if a.number_of_bits == b.number_of_bits { 
                a.literal.partial_cmp(&b.literal).unwrap() 
            } else { 
                b.number_of_bits.partial_cmp(&a.number_of_bits).unwrap() 
            }
        );
        HuffmanTree::new(&literal_order)
    }

    fn weight2bits(max_bits:u8, weight:u8) -> u8 {
        match weight {
            0 => 0,
            w => max_bits + 1 - w,
        }
    }

    pub fn literal_table_to_string(max_bits:u8, src:&Vec<HuffmanEntry>) -> String {
        let table_string = src
            .iter()
            .map(|entry|format!("{}\t{}\t{}",entry.literal, entry.weight, entry.number_of_bits))
            .collect::<Vec<_>>()
            .join("\n");
        format!("max bits:{}\nliteral\tweight\tnumber_of_bits\n{}", max_bits, table_string)
    }

    pub fn root(&self) -> &HuffmanTree {
        &self.huffman_tree
    }
}

impl std::fmt::Debug for HuffmanTable {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("HuffmanTable")
            .field("max_bits", &self.max_bits)
            .finish()
    }
}